"""
集成学习猫狗识别 - Gradio Web应用
加载训练好的集成学习模型，提供图片上传和识别功能
"""

import gradio as gr
import cv2
import numpy as np
import pickle
import os
from PIL import Image

# ==================== 配置参数 ====================
MODEL_PATH = "best_ensemble_model.pkl"
IMAGE_SIZE = (32, 32)  # flat特征使用32x32灰度图

# ==================== 加载模型 ====================
def load_model():
    """加载训练好的集成学习模型"""
    if not os.path.exists(MODEL_PATH):
        raise FileNotFoundError(f"模型文件不存在: {MODEL_PATH}\n请先运行 train_ensemble.py 训练模型")
    
    with open(MODEL_PATH, 'rb') as f:
        model_data = pickle.load(f)
    
    print("=" * 60)
    print("模型加载成功！")
    print("=" * 60)
    print(f"模型名称: {model_data['model_name']}")
    print(f"模型准确率: {model_data['test_accuracy']:.4f} ({model_data['test_accuracy']*100:.2f}%)")
    print(f"特征提取方法: {model_data['feature_method']}")
    print("=" * 60)
    
    return model_data['model'], model_data['model_name'], model_data['test_accuracy']

# 加载模型
try:
    model, model_name, model_accuracy = load_model()
except Exception as e:
    print(f"错误: {e}")
    model, model_name, model_accuracy = None, None, None

# ==================== 图像预处理函数 ====================
def preprocess_image(image):
    """
    将上传的图像预处理为模型输入格式
    
    参数:
        image: PIL Image 或 numpy array
    
    返回:
        processed_image: 预处理后的图像特征向量
    """
    # 转换为numpy数组
    if isinstance(image, Image.Image):
        image = np.array(image)
    
    # 如果是RGB图像，转换为灰度图
    if len(image.shape) == 3:
        image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    
    # 调整大小为32x32
    image = cv2.resize(image, IMAGE_SIZE)
    
    # 展平为1024维向量
    image_flat = image.reshape(1, -1).astype('float32')
    
    # 归一化到[0, 1]
    image_flat = image_flat / 255.0
    
    return image_flat

# ==================== 预测函数 ====================
def predict_cat_dog(image):
    """
    对上传的图像进行猫狗识别
    
    参数:
        image: 上传的图像
    
    返回:
        result_text: 识别结果文本
        confidence_text: 置信度信息
    """
    if model is None:
        return "❌ 错误：模型未加载", "请先运行 train_ensemble.py 训练模型"
    
    try:
        # 预处理图像
        processed_image = preprocess_image(image)
        
        # 模型预测
        prediction = model.predict(processed_image)[0]
        
        # 获取预测概率（如果模型支持）
        if hasattr(model, 'predict_proba'):
            probabilities = model.predict_proba(processed_image)[0]
            confidence = probabilities[prediction]
        else:
            confidence = 1.0  # 如果不支持概率预测，设置为1.0
        
        # 生成结果
        class_name = "🐶 狗 (Dog)" if prediction == 1 else "🐱 猫 (Cat)"
        
        result_text = f"预测结果: {class_name}"
        confidence_text = f"置信度: {confidence:.2%}"
        
        # 添加详细信息
        if hasattr(model, 'predict_proba'):
            cat_prob = probabilities[0]
            dog_prob = probabilities[1]
            detail_text = f"\n\n详细概率:\n• 猫: {cat_prob:.2%}\n• 狗: {dog_prob:.2%}"
        else:
            detail_text = ""
        
        return result_text, confidence_text + detail_text
        
    except Exception as e:
        return f"❌ 预测失败", f"错误信息: {str(e)}"

# ==================== Gradio界面 ====================
def create_interface():
    """创建Gradio Web界面"""
    
    # 自定义CSS样式
    custom_css = """
    .gradio-container {
        font-family: 'Arial', sans-serif;
    }
    .result-box {
        font-size: 24px;
        font-weight: bold;
        padding: 20px;
        border-radius: 10px;
        text-align: center;
    }
    """
    
    # 创建界面
    with gr.Blocks(css=custom_css, title="猫狗识别器") as demo:
        gr.Markdown(
            """
            # 🐱🐶 猫狗识别器
            ### 基于集成学习的图像分类系统
            
            上传一张猫或狗的图片，系统将自动识别并给出预测结果。
            """
        )
        
        # 显示模型信息
        if model is not None:
            gr.Markdown(
                f"""
                **模型信息:**
                - 模型类型: {model_name}
                - 测试准确率: {model_accuracy:.2%}
                - 特征提取: Flat (32×32 灰度图)
                """
            )
        else:
            gr.Markdown(
                """
                ⚠️ **警告:** 模型未加载，请先运行 `train_ensemble.py` 训练模型
                """
            )
        
        with gr.Row():
            with gr.Column():
                # 图片上传组件
                image_input = gr.Image(
                    label="上传图片",
                    type="pil",
                    height=400
                )
                
                # 预测按钮
                predict_btn = gr.Button("🔍 开始识别", variant="primary", size="lg")
                
                # 示例图片（如果有的话）
                gr.Markdown("### 💡 提示")
                gr.Markdown(
                    """
                    - 支持格式: JPG, PNG, JPEG
                    - 建议上传清晰的猫或狗的照片
                    - 识别响应时间 < 2秒
                    """
                )
            
            with gr.Column():
                # 预测结果显示
                result_output = gr.Textbox(
                    label="识别结果",
                    lines=2,
                    interactive=False,
                    elem_classes="result-box"
                )
                
                confidence_output = gr.Textbox(
                    label="置信度信息",
                    lines=5,
                    interactive=False
                )
        
        # 绑定预测函数
        predict_btn.click(
            fn=predict_cat_dog,
            inputs=image_input,
            outputs=[result_output, confidence_output]
        )
        
        # 也可以在上传图片后自动预测
        image_input.change(
            fn=predict_cat_dog,
            inputs=image_input,
            outputs=[result_output, confidence_output]
        )
        
        gr.Markdown(
            """
            ---
            **技术栈:** Python + Scikit-learn + Gradio  
            **开发者:** 集成学习猫狗识别项目
            """
        )
    
    return demo

# ==================== 主函数 ====================
if __name__ == "__main__":
    print("\n" + "=" * 60)
    print("启动 Gradio Web 应用...")
    print("=" * 60)
    
    # 创建并启动界面
    demo = create_interface()
    
    # 启动服务
    demo.launch(
        server_name="127.0.0.1",  # 本地访问
        server_port=7860,  # 端口号
        share=False,  # 不创建公共链接
        show_error=True  # 显示错误信息
    )
